"""
MCP AI模型集成示例
展示如何将AI模型（Qwen）集成到MCP工具中
"""
import asyncio
import json
import os
from datetime import datetime
from typing import Any, Dict, List, Optional

from mcp_models import Tool, ToolParameter, ToolCall, ToolResult
from mcp_server import MCPServer


class QwenModelTool:
    """Qwen模型调用工具"""
    
    def __init__(self, api_key: Optional[str] = None, base_url: str = "https://dashscope.aliyuncs.com/api/v1"):
        """
        初始化Qwen模型工具
        
        Args:
            api_key: 阿里云DashScope API密钥
            base_url: API基础URL
        """
        self.api_key = api_key or os.getenv("DASHSCOPE_API_KEY")
        self.base_url = base_url
        self.model_name = "qwen-turbo"  # 可以根据需要调整模型
        
        if not self.api_key:
            print("⚠️  未设置DASHSCOPE_API_KEY，将使用模拟模式")
    
    async def call_qwen_api(self, messages: List[Dict[str, str]], **kwargs) -> Dict[str, Any]:
        """
        调用Qwen API
        
        Args:
            messages: 对话消息列表
            **kwargs: 其他参数（temperature, max_tokens等）
        """
        if not self.api_key:
            # 模拟模式 - 返回模拟响应
            return await self._mock_response(messages)
        
        try:
            import httpx
            
            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
            }
            
            payload = {
                "model": self.model_name,
                "input": {
                    "messages": messages
                },
                "parameters": {
                    "temperature": kwargs.get("temperature", 0.7),
                    "max_tokens": kwargs.get("max_tokens", 1000),
                    "top_p": kwargs.get("top_p", 0.8),
                }
            }
            
            async with httpx.AsyncClient(timeout=30.0) as client:
                response = await client.post(
                    f"{self.base_url}/services/aigc/text-generation/generation",
                    headers=headers,
                    json=payload
                )
                
                if response.status_code == 200:
                    result = response.json()
                    return {
                        "success": True,
                        "content": result["output"]["text"],
                        "usage": result.get("usage", {}),
                        "model": self.model_name,
                        "timestamp": datetime.now().isoformat()
                    }
                else:
                    return {
                        "success": False,
                        "error": f"API调用失败: {response.status_code} - {response.text}",
                        "timestamp": datetime.now().isoformat()
                    }
        
        except Exception as e:
            return {
                "success": False,
                "error": f"调用异常: {str(e)}",
                "timestamp": datetime.now().isoformat()
            }
    
    async def _mock_response(self, messages: List[Dict[str, str]]) -> Dict[str, Any]:
        """模拟Qwen响应"""
        await asyncio.sleep(0.5)  # 模拟API延迟
        
        user_message = messages[-1]["content"] if messages else "你好"
        
        # 简单的模拟响应逻辑
        if "代码" in user_message or "code" in user_message.lower():
            mock_content = """根据您的请求，我来帮您分析代码：

这段代码的主要功能是实现一个基本的功能。以下是一些建议：

1. **代码结构**: 整体结构清晰，逻辑合理
2. **性能优化**: 可以考虑添加缓存机制
3. **错误处理**: 建议增加更详细的异常处理
4. **代码规范**: 符合Python PEP8规范

如果您需要更具体的建议，请提供详细的代码片段。"""
        
        elif "翻译" in user_message or "translate" in user_message.lower():
            mock_content = """我来帮您翻译内容。请注意以下几点：

1. **准确性**: 保持原文的核心含义
2. **语境**: 考虑上下文语境
3. **专业术语**: 使用准确的专业词汇
4. **语言风格**: 保持适当的语言风格

请提供需要翻译的具体内容，我会为您提供准确的翻译。"""
        
        elif "总结" in user_message or "summary" in user_message.lower():
            mock_content = """基于您的要求，我来提供内容总结：

**主要要点**:
1. 核心概念和关键信息
2. 重要的技术细节
3. 实际应用场景
4. 注意事项和建议

**结论**: 
这是一个综合性的分析，涵盖了主要方面。如需更详细的分析，请提供具体内容。"""
        
        else:
            mock_content = f"""您好！我是Qwen AI助手。

我理解您的问题："{user_message}"

基于这个问题，我的建议是：

1. **分析问题**: 首先需要明确问题的核心
2. **制定方案**: 根据具体情况制定解决方案
3. **实施步骤**: 按步骤执行解决方案
4. **验证结果**: 确保解决方案有效

如果您需要更具体的帮助，请提供更多详细信息，我会为您提供更精准的建议。

*注意：这是模拟响应，请设置DASHSCOPE_API_KEY以使用真实的Qwen API。*"""
        
        return {
            "success": True,
            "content": mock_content,
            "usage": {
                "prompt_tokens": len(user_message),
                "completion_tokens": len(mock_content),
                "total_tokens": len(user_message) + len(mock_content)
            },
            "model": f"{self.model_name} (模拟模式)",
            "timestamp": datetime.now().isoformat(),
            "mock": True
        }


class MCPAIServer(MCPServer):
    """集成AI模型的MCP服务器"""
    
    def __init__(self):
        super().__init__()
        self.qwen_tool = QwenModelTool()
        self._register_ai_tools()
    
    def _register_ai_tools(self):
        """注册AI相关工具"""
        
        # 基础AI对话工具
        self.register_tool(Tool(
            name="ai_chat",
            description="与AI模型进行对话交流",
            parameters=[
                ToolParameter(name="message", type="string", description="用户消息", required=True),
                ToolParameter(name="system_prompt", type="string", description="系统提示词", required=False),
                ToolParameter(name="temperature", type="number", description="生成温度(0-1)", required=False, default=0.7),
                ToolParameter(name="max_tokens", type="integer", description="最大token数", required=False, default=1000)
            ]
        ))
        
        # AI代码审查工具
        self.register_tool(Tool(
            name="ai_code_review",
            description="使用AI进行代码审查",
            parameters=[
                ToolParameter(name="code", type="string", description="要审查的代码", required=True),
                ToolParameter(name="language", type="string", description="编程语言", required=False, default="python"),
                ToolParameter(name="focus", type="string", description="审查重点", required=False, default="安全性、性能、可维护性")
            ]
        ))
        
        # AI文档生成工具
        self.register_tool(Tool(
            name="ai_documentation",
            description="使用AI生成文档",
            parameters=[
                ToolParameter(name="content", type="string", description="要生成文档的内容", required=True),
                ToolParameter(name="doc_type", type="string", description="文档类型", required=False, default="API文档"),
                ToolParameter(name="style", type="string", description="文档风格", required=False, default="详细")
            ]
        ))
        
        # AI内容总结工具
        self.register_tool(Tool(
            name="ai_summarize",
            description="使用AI总结内容",
            parameters=[
                ToolParameter(name="text", type="string", description="要总结的文本", required=True),
                ToolParameter(name="length", type="string", description="总结长度", required=False, default="中等"),
                ToolParameter(name="focus", type="string", description="总结重点", required=False)
            ]
        ))
        
        # AI翻译工具
        self.register_tool(Tool(
            name="ai_translate",
            description="使用AI进行翻译",
            parameters=[
                ToolParameter(name="text", type="string", description="要翻译的文本", required=True),
                ToolParameter(name="source_lang", type="string", description="源语言", required=False, default="auto"),
                ToolParameter(name="target_lang", type="string", description="目标语言", required=True)
            ]
        ))
    
    async def execute_tool(self, tool_call: ToolCall) -> ToolResult:
        """扩展工具执行，添加AI工具支持"""
        
        # 处理AI相关工具
        if tool_call.tool_name in ["ai_chat", "ai_code_review", "ai_documentation", "ai_summarize", "ai_translate"]:
            return await self._execute_ai_tool(tool_call)
        
        # 其他工具使用父类方法
        return await super().execute_tool(tool_call)
    
    async def _execute_ai_tool(self, tool_call: ToolCall) -> ToolResult:
        """执行AI工具"""
        try:
            if tool_call.tool_name == "ai_chat":
                return await self._handle_ai_chat(tool_call.parameters)
            
            elif tool_call.tool_name == "ai_code_review":
                return await self._handle_ai_code_review(tool_call.parameters)
            
            elif tool_call.tool_name == "ai_documentation":
                return await self._handle_ai_documentation(tool_call.parameters)
            
            elif tool_call.tool_name == "ai_summarize":
                return await self._handle_ai_summarize(tool_call.parameters)
            
            elif tool_call.tool_name == "ai_translate":
                return await self._handle_ai_translate(tool_call.parameters)
            
            else:
                return ToolResult(
                    request_id=tool_call.request_id,
                    success=False,
                    error=f"未知的AI工具: {tool_call.tool_name}"
                )
        
        except Exception as e:
            return ToolResult(
                request_id=tool_call.request_id,
                success=False,
                error=f"AI工具执行异常: {str(e)}"
            )
    
    async def _handle_ai_chat(self, params: Dict[str, Any]) -> ToolResult:
        """处理AI对话"""
        message = params.get("message", "")
        system_prompt = params.get("system_prompt")
        temperature = float(params.get("temperature", 0.7))
        max_tokens = int(params.get("max_tokens", 1000))
        
        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": message})
        
        result = await self.qwen_tool.call_qwen_api(
            messages,
            temperature=temperature,
            max_tokens=max_tokens
        )
        
        return ToolResult(
            success=result["success"],
            result=result,
            error=result.get("error")
        )
    
    async def _handle_ai_code_review(self, params: Dict[str, Any]) -> ToolResult:
        """处理AI代码审查"""
        code = params.get("code", "")
        language = params.get("language", "python")
        focus = params.get("focus", "安全性、性能、可维护性")
        
        system_prompt = f"""你是一位经验丰富的{language}代码审查专家。请对以下代码进行全面审查，重点关注：{focus}。

请按以下格式提供审查意见：
1. **代码质量评估**
2. **发现的问题**（按严重程度分类：🔴高危、🟡中等、🔵建议）
3. **改进建议**
4. **最佳实践推荐**

请确保建议具体、可操作。"""
        
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"请审查以下{language}代码：\n\n```{language}\n{code}\n```"}
        ]
        
        result = await self.qwen_tool.call_qwen_api(messages)
        
        return ToolResult(
            success=result["success"],
            result={
                **result,
                "review_type": "code_review",
                "language": language,
                "focus_areas": focus
            },
            error=result.get("error")
        )
    
    async def _handle_ai_documentation(self, params: Dict[str, Any]) -> ToolResult:
        """处理AI文档生成"""
        content = params.get("content", "")
        doc_type = params.get("doc_type", "API文档")
        style = params.get("style", "详细")
        
        system_prompt = f"""你是一位专业的技术文档编写专家。请为以下内容生成{style}的{doc_type}。

文档要求：
1. 结构清晰，层次分明
2. 内容准确，易于理解
3. 包含必要的示例和注意事项
4. 使用Markdown格式
5. 符合{doc_type}的标准规范"""
        
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"请为以下内容生成文档：\n\n{content}"}
        ]
        
        result = await self.qwen_tool.call_qwen_api(messages)
        
        return ToolResult(
            success=result["success"],
            result={
                **result,
                "doc_type": doc_type,
                "style": style
            },
            error=result.get("error")
        )
    
    async def _handle_ai_summarize(self, params: Dict[str, Any]) -> ToolResult:
        """处理AI内容总结"""
        text = params.get("text", "")
        length = params.get("length", "中等")
        focus = params.get("focus")
        
        length_map = {
            "简短": "100字以内",
            "中等": "200-300字",
            "详细": "500字以内"
        }
        
        system_prompt = f"""你是一位专业的内容分析师。请对以下内容进行总结，要求：

1. 总结长度：{length_map.get(length, '200-300字')}
2. 提取关键要点和核心信息
3. 保持原文的主要观点和逻辑
4. 使用清晰简洁的语言"""
        
        if focus:
            system_prompt += f"\n5. 特别关注：{focus}"
        
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"请总结以下内容：\n\n{text}"}
        ]
        
        result = await self.qwen_tool.call_qwen_api(messages)
        
        return ToolResult(
            success=result["success"],
            result={
                **result,
                "summary_length": length,
                "focus_area": focus
            },
            error=result.get("error")
        )
    
    async def _handle_ai_translate(self, params: Dict[str, Any]) -> ToolResult:
        """处理AI翻译"""
        text = params.get("text", "")
        source_lang = params.get("source_lang", "auto")
        target_lang = params.get("target_lang", "")
        
        if not target_lang:
            return ToolResult(
                success=False,
                error="必须指定目标语言"
            )
        
        system_prompt = f"""你是一位专业的翻译专家。请将以下文本从{source_lang}翻译成{target_lang}。

翻译要求：
1. 准确传达原文含义
2. 保持原文的语言风格和语调
3. 使用自然流畅的{target_lang}表达
4. 对专业术语进行准确翻译
5. 保持格式和结构"""
        
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"请翻译：\n\n{text}"}
        ]
        
        result = await self.qwen_tool.call_qwen_api(messages)
        
        return ToolResult(
            success=result["success"],
            result={
                **result,
                "source_language": source_lang,
                "target_language": target_lang
            },
            error=result.get("error")
        )


# 演示用法
async def demo_ai_integration():
    """AI集成演示"""
    print("🤖 MCP AI模型集成演示")
    print("=" * 50)
    
    # 创建AI集成服务器
    ai_server = MCPAIServer()
    
    # 演示AI对话
    print("\n1️⃣ AI对话演示:")
    chat_result = await ai_server._handle_ai_chat({
        "message": "请介绍一下MCP协议的核心概念",
        "temperature": 0.7
    })
    
    if chat_result.success:
        print("✅ AI回复:", chat_result.result["content"][:200] + "...")
    else:
        print("❌ 对话失败:", chat_result.error)
    
    # 演示代码审查
    print("\n2️⃣ AI代码审查演示:")
    code_sample = """
def process_user_input(user_input):
    # 没有输入验证
    result = eval(user_input)  # 安全风险！
    return result
"""
    
    review_result = await ai_server._handle_ai_code_review({
        "code": code_sample,
        "language": "python",
        "focus": "安全性"
    })
    
    if review_result.success:
        print("✅ 代码审查:", review_result.result["content"][:200] + "...")
    else:
        print("❌ 审查失败:", review_result.error)
    
    # 演示文档生成
    print("\n3️⃣ AI文档生成演示:")
    doc_result = await ai_server._handle_ai_documentation({
        "content": "async def send_notification(message: str, priority: str = 'normal') -> Dict[str, Any]",
        "doc_type": "API文档"
    })
    
    if doc_result.success:
        print("✅ 文档生成:", doc_result.result["content"][:200] + "...")
    else:
        print("❌ 生成失败:", doc_result.error)
    
    print(f"\n🔍 API密钥状态: {'已配置' if ai_server.qwen_tool.api_key else '未配置（模拟模式）'}")
    print("💡 提示: 设置环境变量DASHSCOPE_API_KEY以使用真实的Qwen API")


if __name__ == "__main__":
    asyncio.run(demo_ai_integration()) 